import numpy as np
import joblib
import gradio as gr
from PIL import Image
import os
os.environ["LOKY_MAX_CPU_COUNT"] = "12"  # 根据你的系统情况调整
# 加载保存的KNN模型
model = joblib.load('best_knn_model.pkl')
# 加载保存的KNN模型，这样我们可以使用预训练的模型进行预测
def preprocess_image(image):
    image = Image.fromarray(image)
    # 转换为灰度图
    image = image.convert("L")
    # 调整大小为8x8
    image = image.resize((8, 8))
    # 转换为numpy数组并扁平化
    image = np.array(image).flatten()
    # 归一化到0-16的范围
    image = image / 16.0
    return image
# TODO: 定义预测函数，这个函数将用于Gradio接口进行预测
def predict_digit(image):
    # 预处理输入图像
    image = preprocess_image(image)
    image = image.reshape(1, -1)
    prediction = model.predict(image.reshape(1, -1))
    return int(prediction[0])

# TODO: 创建Gradio接口，这个接口将用于用户输入和显示预测结果
iface = gr.Interface(fn=predict_digit,
                     inputs=gr.Sketchpad(shape=(28, 28), label="在此处手写数字"),
                     outputs='label')
# TODO: 启动Gradio接口，用户可以通过这个接口进行交互
iface.launch(share=True)